import os
import time
import gzip
import pickle
import numpy as np
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.utils.data import TensorDataset

from AE import AE


def gpu_setup(use_gpu, gpu_id):
    """GPU设置"""
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)

    if torch.cuda.is_available() and use_gpu:
        print('cuda available with GPU:', torch.cuda.get_device_name(0))
        device = torch.device("cuda")
    else:
        print('cuda not available')
        device = torch.device("cpu")
    return device


def load_data():
    """加载数据集"""
    data_path = "../data/mnist.pkl.gz"

    with gzip.open(data_path, "rb") as f:
        ((x_train, y_train), (x_valid, y_valid), _) = pickle.load(f, encoding="latin-1")

    x_train, y_train, x_valid, y_valid = map(
        torch.tensor, (x_train, y_train, x_valid, y_valid)
    )

    return x_train, y_train, x_valid, y_valid


def train_epoch(model, optimizer, device, data_loader):
    """训练一个epoch"""
    model.train()
    iter = 0
    epoch_loss = 0
    for iter, (x, _) in enumerate(data_loader):
        x = x.to(device)
        x_hat = model.forward(x)
        loss = model.loss(x_hat, x, loss_type='MSE')
        # 反向传播过程
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        epoch_loss += loss.detach().item()
    # epoch_loss /= (iter + 1)
    return epoch_loss, optimizer


def val_epoch(model, device, data_loader):
    """验证一个epoch"""
    model.eval()
    iter = 0
    epoch_loss = 0
    with torch.no_grad():
        for iter, (x, _) in enumerate(data_loader):
            x = x.to(device)
            x_hat = model.forward(x)
            loss = model.loss(x_hat, x, loss_type='MSE')
            epoch_loss += loss.detach().item()
        # epoch_loss /= (iter + 1)
    return epoch_loss


def train_AE(config):
    """训练过程"""
    # 参数设置
    batch_size = config.batch_size
    num_epochs = config.num_epochs
    learning_rate = config.learning_rate
    device = config.device
    # 获取训练集和验证集
    x_train, y_train, x_valid, y_valid = load_data()
    train_dataset = TensorDataset(x_train, y_train)
    val_dataset = TensorDataset(x_valid, y_valid)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)

    model = AE().to(device)
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    # 训练过程
    for epoch in range(num_epochs):
        start = time.time()
        train_loss, optimizer = train_epoch(model, optimizer, device, train_loader)
        val_loss = val_epoch(model, device, val_loader)
        epoch_time = time.time() - start
        print("epoch:{:d}  train_loss:{:.5f}  val_loss:{:.5f}  epoch_time:{:.5f} s".format(
            epoch + 1, train_loss, val_loss, epoch_time))
        AE_test(model, x_valid, device, pause=True)
    return model


def AE_test(model, x_valid, device, index=0, pause=False):
    """输入某个验证集数据, 查看效果"""
    index = index  # 要查看效果的数据下标
    x = x_valid[index].to(device)
    with torch.no_grad():
        x_hat = model.forward(x)
    x = x.detach().cpu().numpy()
    x_hat = x_hat.detach().cpu().numpy()
    fig = plt.figure(0)
    fig.add_subplot(1, 2, 1)
    plt.imshow(x.reshape((28, 28)), cmap="gray")
    fig.add_subplot(1, 2, 2)
    plt.imshow(x_hat.reshape((28, 28)), cmap="gray")
    if pause:
        plt.pause(0.1)
    else:
        plt.show()


def reconstruct_show(model, x_valid, device):
    """输入多个验证集数据进行重构, 查看效果"""
    x_sample = x_valid[10:60]  # 数据范围可自行选择
    x_sample = x_sample.to(device)
    with torch.no_grad():
        x_reconstruct = model.forward(x_sample)
    x_sample = x_sample.detach().cpu().numpy()
    x_reconstruct = x_reconstruct.detach().cpu().numpy()
    plt.figure(figsize=(8, 12))
    for i in range(5):
        plt.subplot(5, 2, 2 * i + 1)
        plt.imshow(x_sample[i].reshape(28, 28), vmin=0, vmax=1, cmap="gray")
        plt.title("Test input")
        plt.colorbar()
        plt.subplot(5, 2, 2 * i + 2)
        plt.imshow(x_reconstruct[i].reshape(28, 28), vmin=0, vmax=1, cmap="gray")
        plt.title("Reconstruction")
        plt.colorbar()
    plt.tight_layout()
    plt.show()

def classifier_show(model, device):
    """查看降维分类效果"""
    x, y, _, _ = load_data()
    x = x.to(device)
    with torch.no_grad():
        latent = model.encoder(x)
    latent = latent.detach().cpu().numpy()
    y = y.detach().cpu().numpy()
    plt.figure(figsize=(8, 6))
    plt.scatter(latent[:, 0], latent[:, 1], c=y)
    plt.colorbar()
    plt.grid()
    plt.show()

class Config(object):
    def __init__(self):
        self.batch_size = 64
        self.num_epochs = 5
        self.learning_rate = 0.001
        self.device = gpu_setup(True, 0)


if __name__ == '__main__':
    # 读取数据集
    x_train, y_train, x_valid, y_valid = load_data()
    # 设置网络参数
    config = Config()
    # 训练模型
    model = train_AE(config)
    # 测试模型
    AE_test(model, x_valid, config.device)
    # 数据重构测试
    reconstruct_show(model, x_valid, config.device)
